package lanQiaoBei.数据结构.并查集;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Stack;

/*
 * 并查集实现
 * */
public class UnionFindSet<V> {
    private class Element {
        public Element(V value) {
            this.value = value;
        }

        private V value;
    }

    //元素值对应的元素表
    private HashMap<V, Element> elementMap;
    //每个元素对应的父
    private HashMap<Element, Element> fatherMap;
    //每个头元素对应集合大小
    private HashMap<Element, Integer> setSizeMap;

    public UnionFindSet(List<V> elements) {
        if (elements == null) return;
        elementMap = new HashMap<>();
        fatherMap = new HashMap<>();
        setSizeMap = new HashMap<>();
        for (V element : elements) {
            Element E = new Element(element);
            elementMap.put(element, E);
            fatherMap.put(E, E);
            setSizeMap.put(E, 1);
        }
    }

    private Element findHead(Element element) {
        Element fatherE = fatherMap.get(element);
        Stack<Element> stack = new Stack<>();
        stack.push(element);
        while (element != fatherE) {
            element = fatherE;
            fatherE = fatherMap.get(element);
            stack.push(element);
        }
        while (!stack.isEmpty()) {
            fatherMap.put(stack.pop(), element);
        }
        return element;
    }

    public boolean isSameSet(V e1, V e2) {
        if (elementMap.containsKey(e1) && elementMap.containsKey(e2)) {
            return findHead(elementMap.get(e1)) == findHead(elementMap.get(e2));
        }
        return false;
    }

    public void union(V e1, V e2) {
        if (elementMap.containsKey(e1) && elementMap.containsKey(e2)) {
            Element head1 = findHead(elementMap.get(e1));
            Element head2 = findHead(elementMap.get(e2));
            if (head1 == head2) return;
            Element big = setSizeMap.get(head1) > setSizeMap.get(head2) ? head1 : head2;
            Element small = big == head1 ? head2 : head1;
            fatherMap.put(small, big);
            setSizeMap.put(big, setSizeMap.get(head1) + setSizeMap.get(head2));
            setSizeMap.remove(small);
        }
    }

    public static void main(String[] args) {
        List<Integer> list = new ArrayList<>();
        for (int i = 0; i < 9; i++) {
            list.add(i);
        }
        UnionFindSet<Integer> unionFindSet = new UnionFindSet<>(list);
        System.out.println(unionFindSet.isSameSet(0, 8));
        unionFindSet.union(0, 1);
        unionFindSet.union(1, 8);
        System.out.println(unionFindSet.isSameSet(0, 8));
    }
}
